#!/usr/bin/perl

use strict;
use warnings;
use FindBin qw($Bin);
use lib "$Bin";

use JSON;
use LWP::Simple;
use LWP::UserAgent;

use MDBUtils;
use MDBAWS;

sub vt_check_file($) {
    my $vt_id = shift;
    init_vt($vt_id);
    my $res = get( "http://www.virustotal.com/analisis/" . $vt_id );
    print "[cached results]\n";
    my $pattern1 = 'Result: <span id="porcentaje"><span.*?>';
    my $pattern2 = '<\/span>';
    if ( $res =~ /$pattern1(\d+)$pattern2\/(\d+) / ) {
        return ( $1, $2 );
    }
    return ( 0, 0 );
}

sub vt_send_file($$) {
    my $filename = shift;
    my $vt_id    = shift;
    my $response;
    my $header;
    my $ua = LWP::UserAgent->new();
    $ua->agent(
        'Mozilla/5.0 (compatible; MSIE 5.5; Windows 98; Win 9x; Windows 2000)'
    );
    my $url =
      "https://www.virustotal.com/vt/en/recepcion?" . get_vt_session_id();
    print "[sending file to VT]\n";
    my $res = $ua->post(
        $url,
        Content_Type => 'form-data',
        Content      => [ archivo => ["$filename"] ]
    );

    $res = "";
    while ( not $res =~ /^\[\"TERMINADO\"/ ) {
        $res = vt_ado( $vt_id, 0, 0 );
        sleep(5);
    }
    my $obj = decode_json($res);

    my $vendors        = $obj->[2];
    my $num_detections = 0;
    my $num_engines    = 0;
    my $vtresult;

    foreach my $vendor (@$vendors) {
        $vtresult->{ "vt_" . $vendor->[0] } = $vendor->[3];
    }

    return ( $num_detections, $num_engines, $vtresult );
}

sub get_vt_id_by_md5($) {
    my $md5 = shift;
    my $ua  = LWP::UserAgent->new;
    my $req =
      HTTP::Request->new(
        POST => "https://www.virustotal.com/vt/en/consultamd5" );
    $req->content_type('application/x-www-form-urlencoded');
    $req->content( 'hash=' . $md5 . '&x=1&y=1' );
    my $res = $ua->request($req);
    if ( $res->{_headers}->{location} =~ /\/resultado.html\?(.*)/ ) {
        return $1;
    }
}

sub init_vt($) {
    my $vt_id = shift;

    return get( "https://www.virustotal.com/vt/en/reanaliza?" . $vt_id );
}

sub vt_ado($$) {
    my $vt_id   = shift;
    my $counter = shift;
    my $command = shift;

    return get( "https://www.virustotal.com/vt/en/resultado?" 
          . $vt_id . "-" 
          . $counter . "-"
          . $command );
}

sub get_vt_session_id() {
    return get("https://www.virustotal.com/vt/en/identificador");
}

my $config = MDBUtils::get_config();
my $queue  = MDBAWS::sqs_get_queue("virustotal");

my $num_detections;
my $num_engines;
my $engine_names;
my $vt_metadata;

while (1) {
    my $msg = MDBAWS::sqs_get_message($queue);

    if ( !defined($msg) ) {
        print "DEBUG: No tickets in the queue, quiting...\n";
        exit(0);
    }
    else {
        print "DEBUG: Picked up SQS ticket: " . $msg->MessageBody() . "\n";

        my $metadata    = MDBAWS::sdb_get_metadata( $msg->MessageBody() );
        my $vt_metadata = {};

        if ( !defined( $metadata->{'virustotal_date'} )
            && defined( $metadata->{'sha256'} ) )
        {

            my $file = MDBAWS::s3_get_file( $config->{'aws'}->{'bucketname'},
                $metadata->{'zipname'} );

            my $vt_id = get_vt_id_by_md5( $metadata->{'md5'} );

            ( $num_detections, $num_engines, $engine_names ) =
              vt_send_file( $file, $vt_id );

            $vt_metadata->{'sha256'} = $metadata->{'sha256'};

            while ( my ( $key, $value ) = each(%$engine_names) ) {
                $vt_metadata->{$key} = $value;
            }
            $vt_metadata->{'virustotal_engines'}  = $num_detections;
            $vt_metadata->{'virustotal_detected'} = $num_detections;
            $vt_metadata->{'virustotal_date'}     = DateTime->now();
            $vt_metadata->{'virustotal_date'}->set_time_zone("GMT");

            MDBAWS::sdb_put_metadata($vt_metadata);

            MDBAWS::sqs_put_queue( "virustotal", $msg->MessageBody() );
            unlink($file);
        }
        print "DEBUG: Deleting SQS ticket: " . $msg->MessageBody() . "\n";
        $queue->DeleteMessage( $msg->ReceiptHandle() );
    }
}
